package com.xlh.dokka.svc.core.model;

import com.jcraft.jsch.JSchException;
import com.xlh.dokka.svc.core.client.DockerClient;
import com.xlh.dokka.svc.properties.DockerProperties;
import com.xlh.dokka.svc.task.ImagePullListener;
import com.xlh.util.SpringContextUtil;
import com.xlh.util.SshUtil;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;

import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;

/**
 * @author cheer
 */
@Data
@Slf4j
public class Node {

    private final String ip;

    private volatile Boolean healthy;

    private ResourceUsage resourceUsage;

    private NodeUsage nodeUsage;

    private SshUtil sshUtil;

    private Boolean sshSuccess = false;

    private Semaphore semaphore = new Semaphore(10);

    private Boolean gpu = false;

    public Node(String ip) {
        this.ip = ip;
    }

    public String createContainerWithGpu(ContainerConfig containerConfig, Integer serialNumber, Boolean retry) {
        DockerProperties dockerProperties = SpringContextUtil.getBean(DockerProperties.class);
        try (SshUtil sshUtil = SshUtil.newInstance(dockerProperties.getSshUsername(), dockerProperties.getSshPassword(), ip,
                dockerProperties.getSshPort())) {

            StringBuilder stringBuilder = new StringBuilder();
            stringBuilder.append("docker run --runtime=nvidia --cap-add=SYS_ADMIN --cap-drop=SYS_BOOT ")
                    .append("-v /sys/fs/cgroup/:/sys/fs/cgroup:ro ");
            if (StringUtils.isNotBlank(containerConfig.getMounts())) {
                stringBuilder.append("-v ").append(containerConfig.getMounts()).append(":/data:rw ");
            }
            stringBuilder.append("-h ").append(containerConfig.getHostname())
                    .append(" --ulimit nofile=65535:65535 ")
                    .append("-p :5901 -p :22 -d ")
                    .append("--cpu-period 100000 --cpu-quota ").append((long) (containerConfig.getCpu() * 100000))
                    .append(" --memory ").append(containerConfig.getMemory())
                    .append(" -e NVIDIA_VISIBLE_DEVICES=").append(serialNumber)
                    .append(" --stop-signal SIGKILL ")
                    .append(containerConfig.getImage())
                    .append(" /usr/sbin/init");

            return sshUtil.exec(stringBuilder.toString()).substring(0, 64);
        } catch (JSchException e) {
            if (retry) {
                return createContainerWithGpu(containerConfig, serialNumber, false);
            }
        } catch (Exception e) {
            ImagePullListener.putTask(ip, containerConfig.getImage());
            log.error("create gpu container failed ", e);
        }
        return null;
    }

    public void initialize() {
        ThreadPoolTaskScheduler taskScheduler = SpringContextUtil.getBean(ThreadPoolTaskScheduler.class);
        taskScheduler.scheduleWithFixedDelay(() -> {
            try {
                syncHealthy();
                syncResourceUsage();
            } catch (Exception e) {
                log.warn("host {} sync health and resource usage exception", ip, e);
            }
        }, 3000);
    }

    public Boolean tryAcquire() {
        try {
            return semaphore.tryAcquire(30, TimeUnit.SECONDS);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            log.warn("thread " + Thread.currentThread().getName() + " is interrupted");
        }
        return false;
    }

    public void releaseAcquire() {
        semaphore.release();
    }

    private void syncHealthy() {
        healthy = DockerClient.ping(ip);
    }

    private void syncResourceUsage() {
        if (healthy && checkSshUtil()) {
            ResourceUsage resourceUsage = new ResourceUsage();
            if (resourceUsage.sync(sshUtil)) {
                this.resourceUsage = resourceUsage;
            }
        }
    }

    private boolean checkSshUtil() {
        if (sshUtil != null && sshUtil.getSession().isConnected()) {
            return true;
        }

        if (sshSuccess) {
            sshSuccess = false;
            log.error("node host {} ssh disconnected", ip);
        }

        DockerProperties dockerProperties = SpringContextUtil.getBean(DockerProperties.class);
        sshUtil = SshUtil.newInstance(dockerProperties.getSshUsername(), dockerProperties.getSshPassword(), ip,
                dockerProperties.getSshPort());

        if (sshUtil.getSession().isConnected()) {
            if (!sshSuccess) {
                sshSuccess = true;
                log.info("node host {} ssh connected", ip);
            }
            return true;
        }
        return false;
    }
}
